#!/usr/bin/env python3
"""
app.py - Flask 后端（含交易明细）
- 记录每笔交易买/卖时间、价格、指标（SMA10/20、High20、RSI）
- 所有策略记录卖出原因（Sell_Reason）
- 新增：提供策略描述（买入/卖出信号及参数），供前端显示
- 前端可在 result.html 中循环渲染 trade_log 表格
- 支持多种波段交易策略：突破、趋势跟随、均值回归、动量、形态、组合策略
- 买卖信号全部在策略函数中计算，backtest仅执行交易
- 优化：权益曲线上标注买点和卖点
"""
import os, uuid, time, traceback
from pathlib import Path
from datetime import datetime as dt

from flask import Flask, render_template, request, redirect, url_for, flash
import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")  # 非 GUI 后端
import matplotlib.pyplot as plt

# 导入策略
from strategy import STRATEGIES, STRATEGY_DESCRIPTIONS

# 设置 HTTP/HTTPS 代理环境变量
os.environ["HTTP_PROXY"] = "http://192.168.1.128:7890"
os.environ["HTTPS_PROXY"] = "http://192.168.1.128:7890"

# ========= 全局路径 ========= #
APP = Flask(__name__)
APP.secret_key = "REPLACE_ME_WITH_RANDOM"
CACHE_DIR = Path("./cache"); CACHE_DIR.mkdir(exist_ok=True)
IMG_DIR = Path("./static/images"); IMG_DIR.mkdir(parents=True, exist_ok=True)

# ========= 常量 ========= #
MAX_RETRY = 3

# ========= 缓存工具 ========= #

def load_cache_if_ok(ticker: str, start: str, end: str, csv_file: Path):
    if not csv_file.exists():
        return None
    df = pd.read_csv(csv_file, header=[0, 1], index_col=0, parse_dates=True)
    if df.empty:
        return None
    if df.index[0] <= pd.to_datetime(start) and df.index[-1] >= pd.to_datetime(end):
        print(f"[INFO] 使用缓存: {csv_file}")
        return df
    return None


def download_data(ticker: str, start: str, end: str) -> pd.DataFrame:
    csv_file = CACHE_DIR / f"{ticker}_{start}_{end}.csv"
    df = load_cache_if_ok(ticker, start, end, csv_file)
    if df is not None:
        return df

    for i in range(1, MAX_RETRY + 1):
        try:
            df = yf.download(
                ticker,
                start=start,
                end=end,
                auto_adjust=True,
                threads=False,
                progress=False,
                group_by="column",
            )
            if not df.empty:
                df.to_csv(csv_file)
                return df
            print(f"[WARN] 第 {i}/{MAX_RETRY} 次下载为空")
        except Exception as e:
            print(f"[WARN] 第 {i}/{MAX_RETRY} 次下载失败: {e}")
        time.sleep(2)
    raise RuntimeError("无法获取数据且无缓存可用")


# ========= 列扁平化 ========= #

def _ensure_close(df: pd.DataFrame) -> pd.DataFrame:
    if "Close" not in df.columns:
        if "Adj Close" in df.columns:
            df = df.rename(columns={"Adj Close": "Close"})
        elif "Price" in df.columns:
            df = df.rename(columns={"Price": "Close"})
        else:
            raise ValueError("缺少 Close 列")
    return df


def flatten_columns(df: pd.DataFrame, ticker: str) -> pd.DataFrame:
    if not isinstance(df.columns, pd.MultiIndex):
        return _ensure_close(df)
    lvl0, lvl1 = df.columns.levels
    if ticker in lvl0:
        df = df[ticker].copy()
    elif ticker in lvl1:
        df = df.xs(ticker, axis=1, level=1).copy()
    else:
        raise ValueError("未找到匹配列")
    return _ensure_close(df)

# ========= 回测函数 ========= #

def backtest(ticker: str, start: str, end: str, cash: float, fee: float = 0.001, strategy_name: str = "breakout"):
    """执行回测并返回统计结果 + 交易明细，仅根据策略信号执行交易"""
    df = download_data(ticker, start, end)
    df = flatten_columns(df, ticker)
    df = df.asfreq("B")
    df["Close"] = df["Close"].ffill()
    if len(df) < 60:
        raise ValueError("数据不足 60 根 Bar")

    df["SMA20"] = df["Close"].rolling(20).mean()
    df["SMA10"] = df["Close"].rolling(10).mean()
    df["SMA50"] = df["Close"].rolling(50).mean()
    df["High20"] = df["Close"].rolling(20).max()
    df["Low"] = df["Low"]
    df["High"] = df["High"]
    pct = df["Close"].pct_change()
    rs = pct.clip(lower=0).rolling(14).mean() / pct.clip(upper=0).abs().rolling(14).mean().replace(0, np.nan)
    df["RSI"] = 100 - 100 / (1 + rs)

    if strategy_name not in STRATEGIES:
        raise ValueError(f"未知策略: {strategy_name}")
    df = STRATEGIES[strategy_name](df)

    pos = 0
    cash_now = cash
    shares = 0
    equity = []
    trade_log = []
    buy_points = []  # 记录买点 (日期, 价格)
    sell_points = []  # 记录卖点 (日期, 价格)

    for date, row in df.iterrows():
        price, sig = row["Close"], row["Signal"]
        if sig == 1 and pos == 0:
            shares = int(cash_now * (1 - fee) // price)
            if shares == 0:
                continue
            cost = shares * price * (1 + fee)
            cash_now -= cost
            pos = 1
            entry = {
                "buy_date": date.strftime("%Y-%m-%d"),
                "buy_price": round(price, 2),
                "buy_SMA10": round(row["SMA10"], 2),
                "buy_SMA20": round(row["SMA20"], 2),
                "buy_High20": round(row["High20"], 2),
                "buy_RSI": round(row["RSI"], 2),
                "shares": shares,
            }
            entry_cost = cost
            buy_points.append((date, price))  # 记录买点
        elif sig == -1 and pos == 1:
            proceeds = shares * price * (1 - fee)
            cash_now += proceeds
            profit = round(proceeds - entry_cost, 2)
            trade_entry = {
                "ID": len(trade_log) + 1,
                **entry,
                "sell_date": date.strftime("%Y-%m-%d"),
                "sell_price": round(price, 2),
                "sell_SMA10": round(row["SMA10"], 2),
                "sell_SMA20": round(row["SMA20"], 2),
                "sell_High20": round(row["High20"], 2),
                "sell_RSI": round(row["RSI"], 2),
                "Profit": profit,
                "sell_reason": row["Sell_Reason"] if pd.notna(row["Sell_Reason"]) else "Unknown",
            }
            trade_log.append(trade_entry)
            shares = 0
            pos = 0
            sell_points.append((date, price))  # 记录卖点
        equity.append((date, cash_now + shares * price))

    if not equity:
        raise ValueError("区间内未交易")

    eq_series = pd.Series([e[1] for e in equity], index=[e[0] for e in equity])
    cagr = (eq_series.iloc[-1] / cash) ** (252 / len(eq_series)) - 1
    maxdd = ((eq_series / eq_series.cummax()) - 1).min()
    wins = [t for t in trade_log if t["Profit"] > 0]
    losses = [t for t in trade_log if t["Profit"] < 0]
    win_rate = len(wins) / len(trade_log) if trade_log else 0
    gross_profit = sum(t["Profit"] for t in wins)
    gross_loss = -sum(t["Profit"] for t in losses)
    profit_factor = gross_profit / gross_loss if gross_loss else np.inf

    # 绘图并标注买点和卖点
    fname = f"{ticker}_{uuid.uuid4().hex[:8]}.png"
    plt.figure(figsize=(10, 6))  # 增大图表尺寸以适应标注
    eq_series.plot()
    plt.title(f"Equity Curve - {ticker}")
    plt.xlabel("Date"); plt.ylabel("Equity ($)")
    plt.grid(True)  # 添加网格线以便查看

    # 标注买点
    buy_dates, buy_prices = zip(*buy_points)
    plt.scatter(buy_dates, [eq_series[date] for date in buy_dates], color='green', marker='^', label='Buy', s=100)
    # 标注卖点
    sell_dates, sell_prices = zip(*sell_points)
    plt.scatter(sell_dates, [eq_series[date] for date in sell_dates], color='red', marker='v', label='Sell', s=100)

    plt.legend()
    plt.tight_layout()
    plt.savefig(IMG_DIR / fname); plt.close()

    return {
        "ticker": ticker,
        "start": start,
        "end": end,
        "cash": cash,
        "trades": len(trade_log),
        "win_rate": win_rate,
        "profit_factor": profit_factor,
        "cagr": cagr,
        "maxdd": maxdd,
        "final_eq": eq_series.iloc[-1],
        "img_file": fname,
        "trade_log": trade_log,
    }

# ========= 路由 ========= #

@APP.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        try:
            ticker = request.form["ticker"].upper().strip()
            start = request.form["start"]
            end = request.form["end"]
            cash = float(request.form["cash"])
            strategy = request.form.get("strategy", "breakout")
            res = backtest(ticker, start, end, cash, strategy_name=strategy)
            return render_template("result.html", **res)
        except Exception as e:
            traceback.print_exc()
            flash(str(e))
            return redirect(url_for("index"))
    return render_template("index.html", strategies=STRATEGIES.keys(), strategy_descriptions=STRATEGY_DESCRIPTIONS)

# ========= 入口 ========= #

if __name__ == "__main__":
    APP.run(host="0.0.0.0", port=8000, debug=True)
